import pickle
import tensorflow as tf
import numpy as np
import os
import sys
import matplotlib.pyplot as plt

FEATURE_PICKLE_PATH = r'_save\trans_learn_tf1x_on_tf2x_1of3_extract_feature.py\v7.0\bottleneck.pickle'
IMG_ROOT_DIR = '../../../../../large_data/CV2/_many_files/flower_photos_liuqilong/'


def sep(label = '', cnt=32):
    print('-' * cnt, label, '-' * cnt, sep='')


tf.random.set_random_seed(1)
np.random.seed(1)

VER = 'v5.0'
ALPHA = 0.001  # learning rate
BATCH_SIZE = 32
N_EPOCHS = 20
FILE_NAME = os.path.basename(__file__)
SAVE_DIR = r'_save\trans_learn_tf1x_2of3_train_model.py\v5.1'

sep('Load features')
with open(FEATURE_PICKLE_PATH, 'br') as f:
    pickle_data = pickle.load(f)

idx2label = pickle_data['idx2label']
label2idx = pickle_data['label2idx']
print(idx2label)
print(label2idx)
N_CLS = len(idx2label.keys())

idx2label = pickle_data['idx2label']

x_test = pickle_data['x_test'].reshape(-1, 2048)
y_test = pickle_data['y_test']
x_path_test = pickle_data['x_path_test']
print('x_test', x_test.shape)
print('y_test', y_test.shape)
print('x_path_test', x_path_test.shape)


def data_loader(x, y, path):
    iters = int(np.ceil(len(x) / BATCH_SIZE))
    for i in range(iters):
        yield x[i*BATCH_SIZE:(i + 1)*BATCH_SIZE], y[i*BATCH_SIZE:(i + 1)*BATCH_SIZE], path[i*BATCH_SIZE:(i + 1)*BATCH_SIZE]


sep('model')
x = tf.placeholder(tf.float32, [None, 2048], name='ph_x')

h = tf.layers.Dense(N_CLS)(x)

sep('Infer')
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())

    saver = tf.train.Saver()
    filename = tf.train.latest_checkpoint(SAVE_DIR)
    saver.restore(sess, filename)

    spr = 4
    spc = 4
    spn = 0
    plt.figure(figsize=[13, 6])
    for i, (bx, by, bpath) in enumerate(data_loader(x_test, y_test, x_path_test)):
        bh = sess.run(h, feed_dict={x: bx})
        bh = np.argmax(bh, axis=1)
        for j, (yi, hi, pathi) in enumerate(zip(by, bh, bpath)):
            spn += 1
            if spn > spr * spc:
                break
            plt.subplot(spr, spc, spn)
            plt.axis('off')
            title = f'{idx2label[yi]}: {idx2label[hi]} ({"V" if yi == hi else "X"})'
            plt.title(title)
            path = os.path.join(IMG_ROOT_DIR, pathi)
            img = plt.imread(path)
            plt.imshow(img)
        if spn > spr * spc:
            break

plt.show()
print('Over')
